package com.smmisia.auth.controller;

import cn.dev33.satoken.stp.StpUtil;
import com.smmisia.common.Result.Result;
import io.github.lnyocly.ai4j.listener.SseListener;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletion;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletionResponse;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatMessage;
import io.github.lnyocly.ai4j.service.IChatService;
import io.github.lnyocly.ai4j.service.PlatformType;
import io.github.lnyocly.ai4j.service.factor.AiService;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RestController
@Tag(name = "openai")
public class OpenAiController {

    // 注入Ai服务
    @Autowired
    private AiService aiService;

    // 使用 Map 存储每个用户的对话历史，键为用户标识，值为对话历史列表
    private Map<Integer, List<ChatMessage>> conversationHistory = new HashMap<>();

    @GetMapping("/chat")
    @Operation(summary = "ai")
    public Result<String> getChatMessage(@RequestParam String question) throws Exception {
        // 获取OpenAi的聊天服务
        IChatService chatService = aiService.getChatService(PlatformType.DEEPSEEK);

        // 创建请求参数
        ChatCompletion chatCompletion = ChatCompletion.builder()
                .model("deepseek-chat")
                .message(ChatMessage.withUser(question))
                .build();


        System.out.println(chatCompletion);

        // 发送chat请求
        ChatCompletionResponse chatCompletionResponse = chatService.chatCompletion(chatCompletion);

        // 获取聊天内容和token消耗
        String content = chatCompletionResponse.getChoices().get(0).getMessage().getContent();
        long totalTokens = chatCompletionResponse.getUsage().getTotalTokens();
        System.out.println("总token消耗: " + totalTokens);

        return Result.ok(content);
    }

    @GetMapping("/chatStream")
    public void getChatMessageStream(@RequestParam String question, HttpServletResponse response) throws Exception {
        // 设置响应的内容类型和字符编码
        response.setContentType("text/event-stream; charset=UTF-8");
        response.setCharacterEncoding("UTF-8");

        Integer userId = StpUtil.getLoginIdAsInt();
        // 获取对话历史
        List<ChatMessage> history = getConversationHistory(userId);


        // 获取OpenAi的聊天服务
        IChatService chatService = aiService.getChatService(PlatformType.DEEPSEEK);

        // 创建请求参数
        ChatCompletion chatCompletion = ChatCompletion.builder()
                .model("deepseek-chat")
                .messages(new ArrayList<>(history)) // 将历史对话添加到 messages 列表
                .message(ChatMessage.withUser(question))
                .stream(true) // 确保开启流式输出
                .build();

        // 获取PrintWriter，并确保使用UTF-8编码
        PrintWriter writer = response.getWriter();

        // 发送chat请求
        SseListener sseListener = new SseListener() {
            @Override
            protected void send() {
                try {
                    // 写入数据时添加换行符，符合SSE协议
                    writer.write("data: " + this.getCurrStr() + "\n\n");
                    writer.flush();
                    System.out.println(this.getCurrStr());
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        };

        // 开始流式响应
        chatService.chatCompletionStream(chatCompletion, sseListener);

        // 不需要手动关闭writer，因为流式响应需要保持连接
        // writer.close(); // 注释掉关闭，保持连接

        // 更新对话历史
        updateConversationHistory(userId, ChatMessage.withUser(question));
        updateConversationHistory(userId, ChatMessage.withAssistant(String.valueOf(sseListener.getOutput())));
    }

    // 获取指定用户的对话历史
    private List<ChatMessage> getConversationHistory(Integer userId) {
        return conversationHistory.computeIfAbsent(userId, k -> new ArrayList<>());
    }

    // 更新指定用户的对话历史
    private void updateConversationHistory(Integer userId, ChatMessage message) {
        List<ChatMessage> history = getConversationHistory(userId);
        history.add(message);
    }
}
